#!/usr/bin/python3

import json
import socket
import sys

QMP_SOCKET = "/tmp/qmp-sock"

class QMPClient:
    def __init__(self, socket_path):
        self.socket_path = socket_path
        self.sock = None

    def connect(self):
        self.sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.sock.connect(self.socket_path)
        # Read the initial QMP greeting
        self._read_response()

    def _send_command(self, cmd):
        self.sock.sendall(json.dumps(cmd).encode('utf-8') + b'\n')

    def _read_response(self):
        response = b""
        while True:
            chunk = self.sock.recv(4096)
            if not chunk:
                break
            response += chunk
            if b'\n' in chunk:
                break
        return json.loads(response.decode('ascii'))

    def execute_command(self, cmd):
        self._send_command(cmd)
        return self._read_response()

    def initialize(self):
        return self.execute_command({"execute": "qmp_capabilities"})

    def query_status(self):
        return self.execute_command({"execute": "query-status"})

    def add_cpu(self, cpu):
        return self.execute_command({
            "execute": "device_add",
            "arguments": {
                "driver": "host-x86_64-cpu",
                "id": f"cpu_{cpu}",
                "core-id": 0,
                "socket-id": cpu,
                "thread-id": 0
            }})

    def add_cpus(self, cpus):
        for c in range(1, int(cpus)):
            print(f'add cpu: {c}')
            response = self.add_cpu(c)
            #print(json.dumps(response, indent=4))
            response = self._read_response()
            #print(json.dumps(response, indent=4))
        return response

    def del_cpus(self, cpus):
        for c in range(1, int(cpus)):
            print(f'del cpu: {c}')
            response = self.del_cpu(c)
        return response

    def del_cpu(self, cpu):
        return self.execute_command({
		    "execute": "device_del",
            "arguments": {
                "id": f"cpu_{cpu}"
			}})

    def query_cpu(self):
        return self.execute_command({"execute": "query-hotpluggable-cpus"})

    def take_snapshot(self, snapshot_name):
        return self.execute_command({
            "execute": "savevm",
            "arguments": {"name": snapshot_name}
        })

    def close(self):
        if self.sock:
            self.sock.close()

def main():
    if len(sys.argv) < 2:
        print("Usage: python qmp_script.py {init|status|qcpu|addcpu <cpu>|delcpu <cpu>|addcpus <num>|delcpus <num>|snapshot <snapshot_name>}")
        sys.exit(1)

    command = sys.argv[1]
    qmp_client = QMPClient(QMP_SOCKET)
    qmp_client.connect()

    qmp_client.execute_command({"execute": "qmp_capabilities"})
    try:
        if command == "init":
            response = qmp_client.initialize()
        elif command == "status":
            response = qmp_client.query_status()
        elif command == "snapshot":
            if len(sys.argv) < 3:
                print("Usage: python qmp_script.py snapshot <snapshot_name>")
                sys.exit(1)
            snapshot_name = sys.argv[2]
            response = qmp_client.take_snapshot(snapshot_name)
        elif command == "qcpu":
            response = qmp_client.query_cpu()
        elif command == "addcpu":
            response = qmp_client.add_cpu(sys.argv[2])
        elif command == "delcpu":
            response = qmp_client.del_cpu(sys.argv[2])
        elif command == "addcpus":
            response = qmp_client.add_cpus(sys.argv[2])
        elif command == "delcpus":
            response = qmp_client.del_cpus(sys.argv[2])
        else:
            print("Unknown command:", command)
            sys.exit(1)

        print(json.dumps(response, indent=4))
    finally:
        qmp_client.close()

if __name__ == "__main__":
    main()
